#/usr/bin/python
# -*- coding: utf-8 -*-
'''
This module contains classes defining a domain and its components:
The class Domain represents a general one-block geometric domain while GlobalDomain 
defines the assembly of several geometric domains. Thus the class Boundary lets 
link two domains together. Class Boundary is associated with class Edge that extracts 
a geometric manifold of a domain (point, curve or surface). 

All these classes are designed in order to facilitate the implementation of domain decomposition of finite
element model by the use of primal, dual, mixte or hybrid approach [1,2]. Moreover
couples of nonconformingly meshed domains may be taken into account by the mortar methods
[3,4] through the use of Mortar Operators defined in the Boundary class. 

Details about classes are given in their own documentation string.   

References:

[1] Gosselet P., Non-overlapping domain decomposition methods in structural mechanics, 
   Rapport interne LMT-Cachan n°265, 2006.
[2] Cresta P.,  
[3] C. Bernardi, Y. Maday, and A. T. Patera, A new non conforming approach to domain 
   decomposition: The mortar element method, in Collège de France Seminar, H. Brezis
   and J.-L. Lions, eds., Pitman, 1994.
[4] Amini A., Analyse multi-échelle de structures hétérogènes par décomposition de domaine: Application aux
   navires à passagers. PhD thesis, University of Nantes, France, 2008. 


Created on 5 mars 2012

@author: Ludovic Barrière and Steven Marguet, Institut Clément Ader, 135 av. de Rangueil, 31077 Toulouse, France

© copyright:
    
    TODO
      
DocTests of the module:
  
  -> see testModule.py

'''
import scipy.sparse as sp
import numpy as np
import copy as cp
import re
from mpi4py import MPI
from itertools import combinations
from collections import namedtuple
#import mpi4py as mpi
from mesh import *

comm = MPI.COMM_WORLD
comm.sp

EDGE = namedtuple("Edge", ["nodeSetUId", "elemSetUId", "type"]) 

class Domain(object):
    '''
    A domain is a geometric entity that is described by one or several 
    discrete counterparts (Meshes). This class enables switching, dialoguing, and so forth, from one mesh
    to another. A domain can enrich, refine or coarse its meshes.
    '''
    
    kwargs = {"meshes": [], "meshesMap": {"U2C": {}, "C2U": []}, "activMesh":None, 
             "meshesTO": {}, "edges": [], "edgesMap": {"U2C": {}, "C2U": []},
             "meshesHLevel": [], "meshesRLevel": [],
            }
    
    def __init__(self, **kwargs):
        '''
        Domain Constructor.
        .meshesMap: Mapping between user and code meshes identifiers
        .meshes: list of meshes of the domain
        .activMesh: current mesh identifier
        .meshesLinks: Projection matrices from one mesh to another.
        
        '''
        
        self.meshes =  kwargs["meshes"]
        self.meshesMap = kwargs["meshesMap"]
        self.activMesh = kwargs["activMesh"]
        self.edges = kwargs["edges"]
        self.edgesMap = kwargs["edgesMap"]
        self.edgeModel = Edge(self)
        
        # transfer operators between meshes 
        self.meshesTO = kwargs["meshesTO"]
        self.meshesHLevel = kwargs["meshesHLevel"]
        self.meshesPLevel = kwargs["meshesPLevel"]
    
    def getLhs(self, meshId, *args):
        pass
        
    def getRhs(self, meshId):
        pass
    
    def setBC(self, meshId):
        pass
    
    def setLC(self, meshId):
        pass
    
    # MULTI-MESH MANAGEMENT:
    
    REMESH_MODES = {
        "default": "Replace",
        "R": "Replace", 
        "AG": "Add Grid",
        }
    
    def hRefineMesh(self, MODE = "default"):
        return NotImplementedError
        
    def pRefineMesh(self, MODE = "default"):
        return NotImplementedError
    
    def hpAdaptMesh(self, MODE = "default"):
        return NotImplementedError
    
    def addMesh(self, mesh, userId):
        codeId = len(self.meshesMap["C2U"])
        self.meshes.append(mesh)
        if userId not in self.meshesMap["C2U"]:
            self.meshesMap["U2C"][userId] = codeId
            self.meshesMap["C2U"].append(userId) 
        else:
            return KeyError("Mesh " + repr(userId) + " already exists.")
        if codeId == 0:
            self.activMesh = 0
        return codeId
    
    def setActivMesh(self, meshCodeId = None, meshUserId = None):
        if meshCodeId:
            self.activMesh = meshCodeId
        if meshUserId:
            self.activMesh = self.meshesMap["U2C"][meshUserId]
    
    def setTransferOperators(self, m1, m2):
        return NotImplementedError
    
    # EDGE DEFINITION:
#    def automaticEdges(self):
#        return NotImplementedError
#    
#    EDGE_KEYWORDS = ("edge", "point", "curve", "surface")
#    
#    
#    def predefinedEdges(self):
#        '''
#        creates edges from predefined node sets in the activated mesh. Node sets names must 
#        contain one of the predefined keywords. 
#        '''
##        mesh = self.meshes[self.activMesh]
##        n_mesh = len(self.meshesMap["C2U"])
##        n_edge = len(self.edgesMap["C2U"])
##        
##        for i,n_set in enumerate(mesh.nodeSetMap["C2U"]):
##            for kw in Domain.EDGE_KEYWORDS:
##                m = re.search(kw, n_set, flags=re.IGNORECASE)
##                if m:
##                    self.edges.append(EDGE_CLASS[kw](mesh.nodeSet[i], self.activMesh, n_mesh))
##                    self.edgesMap["U2C"][n_set] = n_edge
##                    self.edgesMap["C2U"].append(n_set)
##                    n_edge += 1
##                    break
#        
#        return 0
#    
#    EDGES_METHODS = {
#        "predef": predefinedEdges,
#        "auto": automaticEdges,
#        }
    
    def addEdge(self, nodeSetUId=None, elemSetUId=None, nodes=None, elements=None, meshCId = None, edgeUId = None, edgeType = None):
        '''
        Add an edge to the domain. 
        '''
        
        if not meshCId:
            meshCId = self.activMesh
        if nodeSetUId:
            try: 
                self.meshes[meshCId].nodeSetMap["U2C"][nodeSetUId]
            except KeyError:
                if nodes:
                    self.meshes[meshCId].addNodeSet(nodes,nodeSetUId)
                else:
                    return ValueError("Argument missing for defining the edge.")
        else:
            if nodes:
                # create nodeSet from nodes:
                nodeSetUId =  edgeUId + "_edgeNodeSet"
                self.meshes[meshCId].addNodeSet(nodes,nodeSetUId)
            else:
                return ValueError("Argument missing for defining the edge.")        
        if elemSetUId:
            try:
                self.meshes[self.activMesh].elemSetMap["U2C"][elemSetUId]
            except KeyError:
                if elements:
                    self.meshes[meshCId].addElemSet(elements, elemSetUId) 
                else:
                    return ValueError("Argument missing for defining the edge.")
        else:
            if elements:
                # create elementSet from elements:
                elemSetUId =  edgeUId + "_edgeElemSet"
                self.meshes[meshCId].addElemSet(elements, elemSetUId)    
            else:
                return ValueError("Argument missing for defining the edge.")
        
        codeId = len(self.edgesMap["C2U"])
        
        self.edges.append(EDGE(nodeSetUId,elemSetUId,edgeType))
        self.edgesMap["U2C"][edgeUId] = codeId
        self.edgesMap["C2U"].append(edgeUId)
#        self.edges.append(Edge( = 
        
        return codeId
        
    
#    def setEdges(self, METHOD = "predef"):
#        return Domain.EDGES_METHODS[METHOD](self)

class DecoratorDomain(Domain):                
    ''' Minimal dynamic decorator of class Domain.'''
    
    def __init__(self,domain):
        '''Constructor of the decorator. '''
        super.__init__(domain.__dict__())
        self.domain = domain

 
      
class GlobalDomain(DecoratorDomain):
    '''
    The class GlobalDomain() inherits from the class Domain(). It implements a decorator for 
    instances of class Domain(). 
    '''
    
    def __init__(self, domain, subDomains=[], subDomainsUId=[]):
        '''
        GlobalDomain Constructor:
        1. Constitutive domains definition:
        .domains: list of constitutive domains
        .domainsMap: Mapping between user and code identifiers of the domains
        .domainsEdges: lists of edges of each domain
        .domainsLhs: dictionary of Left-Hand-Sides of equilibrium equation(s)
        .domainsRhs: dictionary of Right-Hand-Sides of equilibrium equation(s)
        
        2. Boundaries definition:
        .boundaries: list of boundaries between two domain at a time
        .boundariesMap: Mapping between user and code identifiers of the boundaries
        .boundariesTO: list of trace operators from associated domains towards the boundary
        .boundariesBAO: list of binary assembly operators from domains towards global domain 
        .boundariesSAO: list of signed assembly operators from domains towards global domain
        .boundariesType: list of type of boundaries ( 0 if primal, 1 if dual and 2 if mixed)
        .boundariesConstraints: list of constraints to be applied on d.o.f. of boundaries ( C q = d )
        .boundariesScaling: list of scaling matrices
        
        '''
        # initialize domain:
        kwargs = domain.__dict__()
        super().__init__(kwargs)
        # initialize subdomains:
        self.domains = []
        self.domainsMap = {"U2C": {}, "C2U": []}
#        self.domainsEdges = []
        self.domainsLhs = {}
        self.domainsRhs = {}
        self.domainsBoxes = []
        
        self.boundaries = []
        self.boundariesMap = {"U2C": {}, "C2U": []}
        self.boundariesTO = []
        self.boundariesBAO = []
        self.boundariesSAO = []
        self.boundariesType = []
        self.boundariesConstraints = []
        self.boundariesScaling = []
        
        if len(subDomains) > 0:
            for i, d in enumerate(subDomains):
                self.addDomain(d, subDomainsUId[i])  
        else:
            # domains are predefined in node sets and element sets:
            pass
        
        self.setBoundaries()
            
    # domain methods overload:
    def getLhs(self):
        pass
    
    def getRhs(self):
        pass
    
    def setBC(self):
        pass
    
    def setLC(self):
        pass
    
    def addEdge(self, domainCId, edgeNodes, edgeElements, edgeUId): # domain method overload
        try:
            return self.domains[domainCId].addEdge(nodes=edgeNodes, elements=edgeElements, edgeUId = edgeUId)
        except IndexError:
            return IndexError("Unknown domain identifier.")
        return 1
    
    # own methods:    
    def splitDomain(self):
        pass
        
    def joinDomains(self):
        pass
        
    def addDomain(self, domain, userId):
        codeId = len(self.domainsMap["C2U"])
        self.domains.append(domain)
        if userId not in self.domainsMap["C2U"]:
            self.domainsMap["U2C"][userId] = codeId
            self.domainsMap["C2U"].append(userId)
            self.domainsEdges.append([])  
            self.domainsBoxes.append(self.domainBox(codeId))          
        else:
            return KeyError("Domain " + repr(userId) + " already exists.")
        
    def addBoundary(self, domain1CId, edge1CId, domain2CId=None, edge2CId=None, boundaryUId = None):
        '''
        Add a boundary defined by one domain at least (two domains at most) and corresponding edge(s).
        User identifier of the boundary is either given in input or automatically generated from domains ones.
        '''
        n_boundaries = len(self.boundaries)
        self.boundaries.append(Boundary(domain1CId, domain2CId, edge1CId, edge2CId))
        
        if not boundaryUId:
            boundaryUId = "B_" + str(domain1CId) + "-" + str(domain2CId) 
            while 1:
                v = 0
                try:
                    self.boundariesMap["U2C"][boundaryUId]
                    boundaryUId = boundaryUId + '_' + str(v)
                except KeyError:
                    break
            
        self.boundariesMap["U2C"][boundaryUId] = n_boundaries
        self.boundariesMap["C2U"].append(boundaryUId)
        
    def domainBoxIntersect(self,domain1CId, domain2CId):
        
        I = [[0,0],[0,0],[0,0]]
        
        try:
            I1 = cp.copy(self.domainsBoxes[domain1CId])
        except KeyError:
            I1 = self.domainBox(domain1CId)
        try:
            I2 = cp.copy(self.domainsBoxes[domain2CId])
        except KeyError:
            I2 = self.domainBox(domain2CId)    
        
        for axis in [0,1,2]:
            Iaxis = I1[axis].extend(I2[axis])
            idx = np.argsort(Iaxis)
            if [idx[1],idx[2]] in [[3,0],[2,1]]:
                I[axis]=[Iaxis[idx[1]],Iaxis[idx[2]]]
            else:
                I = None
                break
        return I
    
    def domainIntersect(self, domainCId, nodes=None, elements=None, thirdDomainCId=None, thirdNodes=None, thirdElements=None):
        
        d = self.domains[domainCId]
        m = d.meshes[d.activMesh]
            
        td = self.domains[thirdDomainCId]
        tm = td.meshes[td.activMesh]
        
        edgeNodes = []
        thirdEdgeElements = []
        
        tol = 1.e-15
        
        if not nodes:
            nodes = None
        if not elements:
            elements = None
        
        for nd in nodes:
            X = m.nodesCoo[nd,:]
            
            for el in thirdElements:
                # does X lie on one manifold of the element? 
                N = ELTYPE_TO_MODEL_CLASS[tm.elementsType[el]].interpolation(X, tm.nodesCoo[tm.elementsNodes[el],:], tm.elementsFrame[el])
                flag = True
                for n in N:
                    if n < 0 - tol: # node outside the element
                        flag = False
                        break
                    if n > 1 + tol: # node outside the element
                        flag = False
                        break
                    if n > 0 - tol and n < 0 - tol: # node on one manifold
                        flag = True
                # if yes add nd to edgeNode         
                if flag:
                    edgeNodes.append(nd)
                    thirdEdgeElements.append(el)
                    break
                # if not pass
                
        return edgeNodes, thirdEdgeElements
    
    def domainBox(self,domainCId):
        '''
        Calculates the minimal box containing the domain.
        '''
        
        
        I = [[0,0],[0,0],[0,0]]
        
        d = self.domains[domainCId]
        m = d.meshes[d.activMesh]
        
        for axis in [0,1,2]:
            I[axis][0] = np.min(m.nodesCoo[:,axis], 0)
            I[axis][1] = np.max(m.nodesCoo[:,axis], 0)
        
        return I
    
    BOUNDARIES_METHODS = {
        "auto": automaticBoundaries,
        "predef": predefinedBoundaries,                  
        }
    
    def automaticBoundaries(self):
        '''
        Creates boundaries by domain intersection. 
        '''
#        n_dom = len(self.domains)
#        # compare edges of domains and find neighbours
#        for d1 in xrange(n_dom):
#            edgesd1 = self.domainsEdges[d1] 
#            for d2 in xrange(d1+1, n_dom):
#                edgesd2 = self.domainsEdges[d2]
#                for e2 in edgesd2:
#                    
#                    # edge type of e2:
#                    if isinstance(e2,Point):
#                        edgetype = Point
#                    elif isinstance(e2,Curve):
#                        edgetype = Curve
#                    elif isinstance(e2, Surface):
#                        edgetype = Surface
#                    else:
#                        edgetype = None
#                        
#                    if edgetype:    
#                        for e1 in edgesd1:
#                            # if edge type is the same:
#                            if isinstance(e1,edgetype):
#                                if e1.marker == e2.marker:
#                                    self.boundaries.append(Boundary(d1,d2,e1,e2))
#                                    v = 0
#                                    Bkeys = self.boundariesMap["U2C"].keys()
#                                    n_B = len(Bkeys)
#                                    while 1:
#                                        userId = "B" + repr(d1) + repr(d2) + '.' + repr(v)
#                                        if userId not in Bkeys:
#                                            self.boundariesMap["U2C"][userId] = n_B
#                                            break
#                                        v += 1
#                                    break
        
        # sequential version:
        for d1,d2 in combinations(self.domainsMap["U2C"].values(), 2):
            # interval of intersection
            I = self.domainBoxIntersect(d1, d2)
            if I:
                # selection of nodes and elements located in the interval of intersection calculated above
                nd1, elem1 = self.domainFEEntities(d1, I)
                nd2, elem2 = self.domainFEEntities(d2, I)
                
                # calculation of intersection for each domain 
                
                    # nodes of d1 in d2
                nd1, elem2 = self.domainIntersect(d1, nodes=nd1, elements=elem1, thirdDomain=d2, thirdNodes=nd2, thirdElements=elem2)    
                    # nodes of d2 in d1
                nd2, elem1 = self.domainIntersect(d2, nodes=nd2, elements=elem2, thirdDomain=d1, thirdNodes=nd1, thirdElements=elem1)        
                    # define edges for each domain
                e1 = self.addEdge(d1, nd1, elem1, 'edge_' + str(d1) + '-' + str(d2))
                
                e2 = self.addEdge(d2, nd2, elem2, 'edge_' + str(d2) + '-' + str(d1))   
                # create Boundary instance:
                self.addBoundary(d1, e1, d2, e2)
                              
        return 0
    
    def predefinedBoundaries(self):
        '''
        Creates boundaries by using userId of the edges. user Id of edges must respect the predefined format:
        "edge_" + "userId of domain" + "userId of neighbour domain".
        '''
#        n_dom = len(self.domains)
#        # compare edges of domains and find neighbours
#        for d1 in xrange(n_dom):
#            
#            d1_userId = self.domainsMap["C2U"][d1]
#            d1_edges = self.domainsEdges[d1]
#             
#            for e1 in d1_edges:
#                m = re.search('(edge)_(\w+)-(\w+)', e1, flags = re.IGNORECASE)
#                if m:
#                    d = (int(m.group(1)), int(m.group(2)))
#                    if d1_userId == d[0]:
#                        
#            
#            for d2 in xrange(d1+1, n_dom):
#                edgesd2 = self.domainsEdges[d2]
#                for e2 in edgesd2:
        
        
        return 0
       
    def setBoundaries(self, MODE = "auto"):
        
        GlobalDomain.BOUNDARIES_METHODS[MODE](self)
        
        return 0
    







            
class Boundary(object):
    ''' 
    This class defines a boundary between two domains by linking one edge of a domain to one edge of another.
    It manages the nonconforming meshes by the use of the mortar method. Therefore a virtual boundary is discretized 
    and linked to each constitutive edges.
    --------------. <-> . <-> .--------------
      domain 1    .     
                  . <-> . <->      domain 2
                  . 
    --------------. <-> . <-> .--------------
               Edge1   B12  Edge2
    
    '''
    
    def __init__(self,d1=None,d2=None,e1=None,e2=None):
        '''
        Boundary constructor:
        .neighbours: list of code identifiers of associated domains
        .commonEdges: list of linked edges
        .binaryAssamblyOperators: list of binary assembly operators from 
        
        '''
        self.neighbours = [d1,d2]
        self.commonEdges = [e1,e2]
        #self.edgeTypes
        self.binaryAssemblyOperators = []
        self.signedAssemblyOperators = []
        self.constraints = []
        self.mortarOperators = []
        self.traceOperators = []
        self.interpolation = []

    def enrichBoundary(self):
        pass
    
    def setAssemblyOperators(self):
        pass
        
    
    def setMortarOperators(self):
        pass
    
    

class Edge(object):
    '''
    This class defines an edge (part of the border) of a domain. It is a discrete (n-k)-manifold of the considered 
    domain (which is a n-manifold). An edge is indeed linked to a mesh of the domain from which it is extracted. 
    '''
    
    def __init__(self, domainCId=None, nodeSet=None, elementSet=None, mesh=0, n_mesh=1):
        '''
        Edge constructor:
        .setNode: set of identifiers of the nodes of the domain belonging to the edge.
        .interpolation: interpolation functions of the edge
        .traceOperator: trace operator that extracts d.o.f. of the domain corresponding to the nodes of the edge
        '''
        # related domain id
        self.domain = domainCId
        
        # geometric representation of the edge
        self.nodesCoo = None
        self.simplices = None
        
        # calculated trace operator for the meshes of the domain
        self.traceOperator = [None]*n_mesh
        
        
#        self.nodeSets = [None]*n_mesh
#        self.nodeSets[mesh] = nodeSet
#        
#        self.elements = [None]*n_mesh
#        self.elements[mesh] = elementSet
        
#        self.marker = None
        
    def calcMarker(self, domain=None, nodeSet=None, elementSet=None, mesh=0, n_mesh=1 ):
        '''
        Returns a marker for the edge in order to compare it with other.
        '''
        pass
    
    def getTraceOperator(self, domain=None, nodeSet=None, elementSet=None, mesh=0, n_mesh=1):
        
        pass
    
    def getInterpolation(self, domain=None, nodeSet=None, elementSet=None, mesh=0, n_mesh=1):
        
        pass
    
    
    
class Point(Edge):
    '''
    Defines an edge as a point  
    '''
    
    def __init__(self, edge, node, mesh = 0, n_mesh = 1):
        '''
        
        '''
        kwargs = edge.__dict__()
        super().__init__(kwargs)
        self.calcMarker()
        
    
    def calcMarker(self):
        self.marker = self.nodes[0]
    
class Curve(Edge):
    '''
    Defines a curve as an edge
    '''
    
    def __init__(self, edge, nodeset, mesh = 0, n_mesh = 1):
        '''
        
        '''
        edge.__init__(self, nodeset, mesh, n_mesh)
        
    
    def calcMarker(self):
        pass
    
    
class Surface(Edge):
    '''
    Defines a surface as an edge
    '''
    
    def __init__(self, nodeset, mesh = 0, n_mesh = 1):
        '''
        
        '''
        Edge.__init__(self, nodeset, mesh, n_mesh)
    
    def calcMarker(self):
        pass
    
EDGE_CLASS = {
        "edge": Edge,
        "point": Point,          
        "curve": Curve,
        "surface": Surface,
        }           